import numpy as np
from sklearn.datasets import load_digits
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt


if __name__ == '__main__':
    data, labels = load_digits(return_X_y=True)
    (n_samples, n_features), n_digits = data.shape, np.unique(labels).size

    m = 31
    inertias = []
    for n_clusters in range(1, m):
        kmeans = KMeans(init="k-means++", n_clusters=n_clusters, n_init=4)
        kmeans.fit(data)
        inertias.append(kmeans.inertia_)

    plt.plot(list(range(1, m)), inertias)
    plt.plot(10, inertias[9], 'r+')
    plt.show()
